package abacab.serialization.huffman;

import java.io.IOException;
import java.io.Reader;
import java.io.Writer;
import java.util.HashMap;
import java.util.Map;
import java.util.NavigableSet;
import java.util.Stack;
import java.util.TreeSet;

/**
 * Encodes with Huffman algorithm
 */
public final class Huffman {
	/**
	 * Get the frequencies of the characters
	 * 
	 * @param reader
	 *          the reader
	 * @return the list of characters ordered by frequency
	 * @throws IOException
	 *           if an error occurred while reading
	 */
	public NavigableSet<Node> getFrequency(Reader reader) throws IOException {
		NavigableSet<Node> nodes = new TreeSet<Node>(Node.COMPARATOR);
		Map<Character, LeafNode> chars = new HashMap<Character, LeafNode>();
		int c;
		while ((c = reader.read()) >= 0) {
			Character ch = Character.valueOf((char) c);
			LeafNode node = chars.get(ch);
			if (node == null) {
				node = new LeafNode(ch.charValue(), 1);
				chars.put(ch, node);
			}
			else {
				node.addWeight();
			}
		}
		for (LeafNode node : chars.values()) {
			nodes.add(node);
		}
		nodes.add(new LeafNode('\0', 1));
		return nodes;
	}

	/**
	 * Generate the binary tree. The provided list will be altered in the process.
	 * 
	 * @param nodes
	 *          the nodes
	 * @return the root of the tree
	 * @throws IllegalArgumentException
	 *           if the provided list of nodes is NULL or empty
	 */
	public Node generateBinaryTree(NavigableSet<Node> nodes) throws IllegalArgumentException {
		if (nodes == null || nodes.isEmpty()) {
			throw new IllegalArgumentException("You must provide at least 1 node to build a binary tree");
		}
		if (nodes.size() == 1) {
			return nodes.pollFirst();
		}
		while (true) {
			Node n1 = nodes.pollFirst();
			Node n2 = nodes.pollFirst();
			Node n3 = new IntermediaryNode(n1, n2);
			if (nodes.isEmpty()) {
				return n3;
			}
			nodes.add(n3);
		}
	}

	/**
	 * Generate the coding table
	 * 
	 * @param root
	 *          the root
	 * @return the coding table
	 */
	public Map<Character, Code> generateCodingTable(Node root) {
		Map<Character, Code> codingTable = new HashMap<Character, Code>();
		Stack<Node> stack = new Stack<Node>();
		Node n = root;
		NavigationState state = NavigationState.COMMING_DOWN;
		int mask = 0;
		int maskSize = 0;
		do {
			switch (state) {
				case COMMING_DOWN:
					if (n.getLeft() != null) {
						stack.push(n);
						n = n.getLeft();
						mask <<= 1;
						maskSize++;
					}
					else {
						codingTable.put(Character.valueOf(n.getCharacter()), new Code(n.getCharacter(), mask, maskSize));
						n = stack.pop();
						state = NavigationState.COMMING_BACK;
					}
					break;
				case COMMING_BACK:
					if ((mask & 1) == 0) {
						// Coming back from left
						stack.push(n);
						n = n.getRight();
						mask |= 1;
						state = NavigationState.COMMING_DOWN;
					}
					else {
						n = stack.pop();
						mask >>= 1;
						maskSize--;
					}
					break;
			}
		} while (n != root || (mask & 1) == 0);
		return codingTable;
	}

	/**
	 * Encode the text
	 * 
	 * @param reader
	 *          the reader
	 * @param context
	 *          the encoding context
	 * @param codingTable
	 *          the coding table
	 * @throws IOException
	 *           if an error occurred while reading or writing
	 */
	public void encode(Reader reader, EncodingContext context, Map<Character, Code> codingTable) throws IOException {
		int c;
		while ((c = reader.read()) >= 0) {
			_encodeChar((char) c, context, codingTable);
		}
		_encodeChar('\0', context, codingTable);
	}

	/**
	 * Encode a character into the stream
	 * 
	 * @param c
	 *          the character to encode
	 * @param context
	 *          the encoding context
	 * @param codingTable
	 *          the coding table
	 * @throws IOException
	 *           if an error occurred while reading or writing
	 */
	private final void _encodeChar(char c, EncodingContext context, Map<Character, Code> codingTable) throws IOException {
		Character ch = Character.valueOf(c);
		Code code = codingTable.get(ch);
		int maskSize = code.getMaskSize();
		maskSize -= context.getBitsRemaining();
		while (maskSize > 0) {
			context.updateCurrentByte(code.getCode() >> maskSize);
			context.commitCurrentByte();
			maskSize -= context.getBitsRemaining();
		}
		if (maskSize < 0) {
			context.updateCurrentByte(code.getCode() << -maskSize);
			context.setBitsRemaining(-maskSize);
		}
		else {
			context.updateCurrentByte(code.getCode());
			context.commitCurrentByte();
		}
	}

	/**
	 * Decode the text
	 * 
	 * @param writer
	 *          the writer
	 * @param context
	 *          the decoding context
	 * @param root
	 *          the root of the binary tree (dictionary)
	 * @throws IOException
	 */
	public void decode(Writer writer, DecodingContext context, Node root) throws IOException {
		while (true) {
			Node n = root;
			do {
				n = context.isNextBitSet() ? n.getRight() : n.getLeft();
			} while (n.getLeft() != null);
			char c = n.getCharacter();
			if (c == '\0') {
				break;
			}
			writer.write(c);
		}
	}

	/**
	 * The navigation state
	 */
	private static enum NavigationState {
		/** Coming down the nodes */
		COMMING_DOWN,
		/** Coming back */
		COMMING_BACK
	}
}